import os
import matplotlib.pyplot as plt
import mindspore.dataset as ds


def cifar(dataset_dir, which_cifar):
    if which_cifar == "cifar10":
        # mindspore built-in api
        dataset = ds.Cifar10Dataset(dataset_dir=dataset_dir,              # dir hold mnist dataset
                                    usage="train",                        # "train"/"val"/"all" part of dataset
                                    num_parallel_workers=4,               # multithreading
                                    sampler=ds.RandomSampler(False, 10))  # use random sampler to sample 10 images
    elif which_cifar == "cifar100":
        # mindspore built-in api
        dataset = ds.Cifar100Dataset(dataset_dir=dataset_dir,              # dir hold mnist dataset
                                     usage="val",                          # "train"/"val"/"all" part of dataset
                                     num_parallel_workers=4,               # multithreading
                                     sampler=ds.SequentialSampler(0, 10))  # use sequential sampler to sample 10 images
    else:
        raise Exception("not support cifar type")

    # To see what samples do we get from dataset
    for index, data in enumerate(dataset.create_dict_iterator(output_numpy=True)):
        label_name = "label" if which_cifar == "cifar10" else "fine_label"
        print(data["image"].shape, data[label_name])
        plt.subplot(2, 5, index+1)
        plt.imshow(data["image"])
        plt.title(data[label_name])
    plt.show()


def fetch_dataset(dataset_dir, which_cifar, unzip=False, download=False):
    if which_cifar not in ["cifar10", "cifar100"]:
        raise Exception("only support cifar10 or cifar100")
    if not os.path.exists(dataset_dir):
        os.makedirs(dataset_dir)

    if which_cifar == "cifar10":
        if download:
            os.system("wget -P " + dataset_dir + " " + "http://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz")
            unzip = True
        if unzip:
            if os.path.exists(dataset_dir + "/cifar-10-binary.tar.gz"):
                os.system("tar -zxvf " + dataset_dir + "/cifar-10-binary.tar.gz -C " + dataset_dir)
        if len(os.listdir(dataset_dir + "/cifar-10-batches-bin")) != 8:
            raise Exception("Cifar10 dataset is broken, the nums of source file should be 8")

    if which_cifar == "cifar100":
        if download:
            os.system("wget -P " + dataset_dir + " " + "http://www.cs.toronto.edu/~kriz/cifar-100-binary.tar.gz")
            unzip = True
        if unzip:
            if os.path.exists(dataset_dir + "/cifar-100-binary.tar.gz"):
                os.system("tar -zxvf " + dataset_dir + "/cifar-100-binary.tar.gz -C " + dataset_dir)
        if len(os.listdir(dataset_dir + "/cifar-100-binary")) != 5:
            raise Exception("Cifar100 dataset is broken, the nums of source file should be 4")



if __name__ == '__main__':
    dataset_dir = "./dataset/cifar10"
    fetch_dataset(dataset_dir=dataset_dir, which_cifar="cifar10", unzip=False, download=False)
    cifar(dataset_dir + "/cifar-10-batches-bin", "cifar10")

    dataset_dir = "./dataset/cifar100"
    fetch_dataset(dataset_dir=dataset_dir, which_cifar="cifar100", unzip=False, download=False)
    cifar(dataset_dir + "/cifar-100-binary", "cifar100")
